# 人工智能NLP-Agent数字人项目-04-基金数据问答任务工单V1.1-20250214
import abc
from typing import Any
from langchain.tools import BaseTool
import logging
import sys
import pandas as pd
import sqlite3
from utils.configFinRAG import sql_examples_path
from FinSQL_01_generate import generate_sql
from FinSQL_02_query import query_db
from FinSQL_03_answer_from_SQL import generate_answer
from utils.instances import TOKENIZER, LLM

# 配置日志
logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

# 全局变量，用于缓存示例数据
g_example_question_list = []
g_example_sql_list = []
g_example_fa_list = []
g_example_info_list = []
g_example_token_list = []


def load_sql_examples():
    """
    加载SQL示例数据并缓存到全局变量中。
    """
    global g_example_question_list, g_example_sql_list, g_example_fa_list, g_example_info_list, g_example_token_list
    if not g_example_question_list:  # 避免重复加载
        sql_examples_file = pd.read_csv(sql_examples_path, delimiter=",", header=0)
        for index, row in sql_examples_file.iterrows():
            g_example_question_list.append(row['问题'])
            g_example_sql_list.append(row['SQL'])
            g_example_info_list.append(row['资料'])
            g_example_fa_list.append(row['FA'])
            tokens = TOKENIZER(row['问题'])
            g_example_token_list.append(tokens['input_ids'])


class FinSQLRAG(BaseTool, abc.ABC):
    name = "查询金融数据库"
    description = "当被问到金融查询相关的问题时，会去金融数据库检索结果"

    def __init__(self):
        super().__init__()

    async def _arun(self, *args: Any, **kwargs: Any) -> Any:
        # 用例中没有用到 arun 不予具体实现
        pass

    def _run(self, query: str) -> str:
        """
        执行金融数据库查询任务。
        :param query: 用户输入的查询问题
        :return: 查询结果或错误信息
        """
        try:
            # 加载SQL示例数据
            load_sql_examples()

            # 生成SQL查询语句
            result_prompt, sql = generate_sql(query, LLM, g_example_question_list, g_example_sql_list, g_example_token_list)
            logging.info(f"生成的SQL查询语句: {sql}")

            # 连接数据库并执行查询
            db_path = '/Users/wanghr/Documents/八维研修/项目/999-项目/fay数字人/金融场景智能问答系统/bs_challenge_financial_14b_dataset/dataset/博金杯比赛数据.db'
            conn = sqlite3.connect(db_path)
            cs = conn.cursor()
            success_flag, exc_result = query_db(sql, cs)
            conn.close()

            if not success_flag:
                return "数据库查询失败，请检查SQL语句或数据库连接。"

            # 生成自然语言回答
            answer = generate_answer(query, exc_result, LLM, g_example_question_list, g_example_info_list, g_example_fa_list, g_example_token_list)
            return answer

        except Exception as e:
            logging.error(f"FinSQLRAG处理异常: {e}", exc_info=True)
            return "FinSQLRAG处理异常，请稍后重试或联系管理员。"


if __name__ == "__main__":
    # 示例用法
    tool = FinSQLRAG()
    result = tool.run("请帮我计算，在20210715，中信行业分类划分的一级行业为消费者服务行业中，涨跌幅最大股票的股票代码是？涨跌幅是多少？百分数保留两位小数。股票涨跌幅定义为：（收盘价 - 前一日收盘价 / 前一日收盘价）* 100%。")
    print(result)